package main

import (
	"context"
	"github.com/elazarl/goproxy"
	"golang.org/x/crypto/ssh"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"ssh-tunnel/util"
	"sync"
	"sync/atomic"
	"time"
)

type ReverseHttpProxy struct {
	Server            string
	User              string
	Password          string
	Timeout           int
	RetryInterval     int
	KeepAliveInterval int
	KeepAliveCountMax int
	RemoteBindAddr    string
	HttpUnixAddr      string

	ctx context.Context
}

func startReverseHttpServer() {
	r := &ReverseHttpProxy{
		Server:            *server,
		User:              *user,
		Password:          *password,
		Timeout:           *timeout,
		RetryInterval:     *retry,
		KeepAliveInterval: 30,
		KeepAliveCountMax: 10,
		HttpUnixAddr:      "ReverseHttpProxy.sock",
		RemoteBindAddr:    *httpProxyRemoteBindAddr,
	}
	ctx, cancelFunc := context.WithCancel(context.Background())
	r.ctx = ctx

	registerOnShutdown(func() {
		cancelFunc()
	})

	go r.sshClient()

	go r.httpProxy()
}

func needOpenReverseHttpServer() bool {
	if len(*server) == 0 {
		return false
	}
	if len(*httpProxyRemoteBindAddr) == 0 {
		return false
	}
	return true
}

func (r *ReverseHttpProxy) httpProxy() error {
	os.Remove(r.HttpUnixAddr)

	proxy := goproxy.NewProxyHttpServer()
	proxy.Verbose = false
	unixAddr, err := net.ResolveUnixAddr("unix", r.HttpUnixAddr)
	if err != nil {
		return err
	}
	ln, err := net.ListenUnix("unix", unixAddr)
	if err != nil {
		return err
	}
	registerOnShutdown(func() {
		log.Printf("[http proxy] close %s\n", unixAddr.String())
		ln.Close()
		os.Remove(r.HttpUnixAddr)
	})
	log.Printf("[http proxy] http proxy listening at %s\n", unixAddr.String())
	return http.Serve(ln, proxy)
}

func (r *ReverseHttpProxy) sshClient() {
	for {
		select {
		case <-r.ctx.Done():
			return
		default:
			r.sshClient0()
			<-time.After(time.Duration(r.RetryInterval) * time.Second)
			log.Printf("[http proxy] ssh connection retrying...\n")
		}
	}
}

func (r *ReverseHttpProxy) sshClient0() {
	var once sync.Once // Only print errors once per session
	config := &ssh.ClientConfig{
		Timeout: time.Duration(r.Timeout) * time.Second,
		User:    r.User,
		Auth: []ssh.AuthMethod{
			ssh.Password(r.Password),
			ssh.PasswordCallback(func() (string, error) {
				return "", nil
			}),
			ssh.KeyboardInteractive(func(user, instruction string, questions []string, echos []bool) (answers []string, err error) {
				log.Println(`user: ` + user)
				log.Println(`Instruction: ` + instruction)
				log.Println(`Questions:`)
				for q := range questions {
					log.Printf(`%d\n`, q)
				}
				countQuestions := len(questions)
				if countQuestions == 1 {
					answers = make([]string, countQuestions, countQuestions)
					answers[0] = ""
				} else if countQuestions > 1 {
					log.Println(`The ssh server is asking multiple questions! s program cannot handle s case.`)
				}
				err = nil
				return
			}),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Config: ssh.Config{
			Ciphers: []string{
				`aes128-cbc`,
				`aes128-ctr`,
				`aes192-ctr`,
				`aes256-ctr`,
				`aes128-gcm@openssh.com`,
				`arcfour256`,
				`arcfour128`,
			},
			MACs: []string{`hmac-sha1`},
		},
	}
	log.Printf("[http proxy] ssh connecting %s...\n", r.Server)
	// Connect to the server host via SSH.
	cl, err := ssh.Dial("tcp", r.Server, config)
	if err != nil {
		once.Do(func() { log.Printf("[http proxy] ssh dial error: %v", err) })
		return
	}
	log.Printf("[http proxy] ssh connected %s\n", r.Server)
	registerOnShutdown(func() {
		log.Printf("[http proxy] ssh disconnect %s...\n", r.Server)
		cl.Close()
	})
	go r.keepAliveMonitor(&once, cl)
	defer cl.Close()
	ln, err := cl.Listen("tcp", r.RemoteBindAddr)
	if err != nil {
		once.Do(func() { log.Printf("[http proxy] bind error: %v", err) })
		return
	}
	// The socket is binded. Make sure we close it eventually.
	bindCtx, cancel := context.WithCancel(r.ctx)
	defer cancel()
	go func() {
		cl.Wait()
		cancel()
	}()
	go func() {
		<-bindCtx.Done()
		once.Do(func() {}) // Suppress future errors
		ln.Close()
	}()
	// Accept all incoming connections.
	for {
		remoteConn, err := ln.Accept()
		if err != nil {
			once.Do(func() {
				log.Printf("[http proxy] accept error: %v", err)
			})
			return
		}
		go r.dialTunnel(bindCtx, remoteConn)
	}
}

func (r *ReverseHttpProxy) keepAliveMonitor(once *sync.Once, client *ssh.Client) {
	if r.KeepAliveInterval == 0 || r.KeepAliveCountMax == 0 {
		return
	}

	// Detect when the ssh connection is closed.
	wait := make(chan error, 1)
	go func() {
		wait <- client.Wait()
	}()

	// Repeatedly check if the remote server is still alive.
	var aliveCount int32
	ticker := time.NewTicker(time.Duration(r.KeepAliveInterval) * time.Second)
	defer ticker.Stop()
	for {
		select {
		case err := <-wait:
			if err != nil && err != io.EOF {
				once.Do(func() { log.Printf("[http proxy] ssh error: %v", err) })
			}
			return
		case <-ticker.C:
			if n := atomic.AddInt32(&aliveCount, 1); n > int32(r.KeepAliveCountMax) {
				once.Do(func() { log.Printf("[http proxy] ssh keep-alive termination") })
				client.Close()
				return
			}
		}

		go func() {
			_, _, err := client.SendRequest("keepalive@openssh.com", true, nil)
			if err == nil {
				atomic.StoreInt32(&aliveCount, 0)
			}
		}()
	}
}

func (r *ReverseHttpProxy) dialTunnel(ctx context.Context, remoteConn net.Conn) {
	connCtx, cancel := context.WithCancel(ctx)
	defer cancel()
	go func() {
		<-connCtx.Done()
		remoteConn.Close()
	}()
	localConn, err := net.Dial("unix", r.HttpUnixAddr)
	if err != nil {
		log.Printf("[http proxy] dial error: %v", err)
		return
	}
	go func() {
		<-connCtx.Done()
		localConn.Close()
	}()
	log.Printf("[http proxy] Connection established R:%s <==> L:%s", remoteConn.RemoteAddr(), localConn.RemoteAddr())
	defer log.Printf("[http proxy] Connection closed R:%s <==> L:%s", remoteConn.RemoteAddr(), localConn.RemoteAddr())
	util.IOTransport(remoteConn, localConn)
}
